package com.lengxf.flink.sink;

import com.lengxf.flink.bean.WaterSensor;
import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy;
import org.apache.flink.connector.datagen.source.DataGeneratorSource;
import org.apache.flink.connector.datagen.source.GeneratorFunction;
import org.apache.flink.connector.jdbc.JdbcConnectionOptions;
import org.apache.flink.connector.jdbc.JdbcExecutionOptions;
import org.apache.flink.connector.jdbc.JdbcSink;
import org.apache.flink.connector.jdbc.JdbcStatementBuilder;
import org.apache.flink.streaming.api.CheckpointingMode;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;

/**
 * 消息写入mysql
 *
 * @author Lengxf
 */
public class SinkMysql {

    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(2);
        env.enableCheckpointing(2000, CheckpointingMode.EXACTLY_ONCE);
        DataGeneratorSource<WaterSensor> source = new DataGeneratorSource<>(
                (GeneratorFunction<Long, WaterSensor>) aLong -> new WaterSensor(String.valueOf(aLong), Math.toIntExact(aLong), aLong),
                100,
                RateLimiterStrategy.perSecond(1),
                Types.POJO(WaterSensor.class)
        );
        //数据源消息
        DataStreamSource<WaterSensor> waterSensorSource = env.fromSource(source, WatermarkStrategy.noWatermarks(), "data-generator");


        SinkFunction<WaterSensor> sink = JdbcSink.sink(getSql(),
                getWaterSensorJdbcStatementBuilder(),
                getJdbcExecutionOptions(),
                getJdbc()
        );
        waterSensorSource.addSink(sink);
        waterSensorSource.print();
        env.execute();
    }

    private static JdbcConnectionOptions getJdbc() {
        return new JdbcConnectionOptions.JdbcConnectionOptionsBuilder()
                .withUrl("jdbc:mysql://127.0.0.1:3306/test_1?useSSL=false&allowPublicKeyRetrieval=true")
                .withDriverName("com.mysql.cj.jdbc.Driver")
                .withUsername("root")
                .withPassword("123456")
                .withConnectionCheckTimeoutSeconds(10)
                .build();
    }

    private static JdbcExecutionOptions getJdbcExecutionOptions() {
        return JdbcExecutionOptions.builder()
                .withBatchIntervalMs(3)
                .withBatchSize(100)
                .withBatchIntervalMs(3000)
                .build();
    }

    private static String getSql() {
        return "INSERT INTO `test_1`.`water_sensor` (`id`, `vc`, `ts`) VALUES (?, ?, ?)";
    }

    private static JdbcStatementBuilder<WaterSensor> getWaterSensorJdbcStatementBuilder() {
        return (preparedStatement, waterSensor)
                -> {
            preparedStatement.setString(1, waterSensor.id);
            preparedStatement.setInt(2, waterSensor.vc);
            preparedStatement.setLong(3, waterSensor.ts);
        };
    }

}
